package com.joysuch.wwyt.util;

import lombok.extern.slf4j.Slf4j;
import org.codehaus.groovy.jsr223.GroovyScriptEngineFactory;

import javax.script.*;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;
import java.util.function.Supplier;

/**
 * groovy脚本执行引擎帮助类
 * 注意：已经编译过的脚本不会释放，通过identity和version拼接后作为唯一字段
 *
 * @author zhaoq
 * 2021-08-06
 */
@Slf4j
public class CachedGroovyScriptEngineUtil {

    private static GroovyScriptEngineFactory factory = new GroovyScriptEngineFactory();
    private static ConcurrentHashMap<String, ScriptEngineWrapper> engineMap = new ConcurrentHashMap<>();
    private static AtomicLong compileCount = new AtomicLong();

    private static class ScriptEngineWrapper {
        private ScriptEngine engine;
        private String code;

        public ScriptEngineWrapper(ScriptEngine engine, String code) {
            this.engine = engine;
            this.code = code;
        }

        public ScriptEngine getEngine() {
            return engine;
        }

        public String getCode() {
            return code;
        }
    }
    /**
     * 执行一个groovy脚本
     * 通过脚本md5唯一标识一个脚本，
     * 多次调用时如果与之前的md5相同，则使用缓存中的脚本引擎
     *
     * @param code             脚本代码
     * @param bindingsSupplier 提供执行脚本需要的Bindings对象，可为空
     * @param func             要调用的脚本中函数名称
     * @param args             要调用的脚本中函数参数
     * @param <T>              脚本中函数返回类型
     * @return 脚本调用返回结果
     * @throws ScriptException       脚本执行异常
     * @throws NoSuchMethodException 脚本中没有指定函数异常
     */
    public static <T> T runScript(String code, Supplier<Bindings> bindingsSupplier, String func, Object... args) throws ScriptException, NoSuchMethodException {
        String key = Md5Util.md5(code);
        ScriptEngineWrapper engineWrapper = engineMap.computeIfAbsent(key, (k) -> {
            compileCount.incrementAndGet();
            ScriptEngine scriptEngine = factory.getScriptEngine();
            Bindings bindings = null;
            if (bindingsSupplier != null) {
                bindings = bindingsSupplier.get();
            }
            if (bindings == null) {
                bindings = new SimpleBindings();
            }
            try {
                scriptEngine.eval(code, bindings);
                log.info("groovy script compiled, total compile count {} code {}", compileCount.get(), code);
            } catch (ScriptException e) {
                log.error("groovy script exception, code " + code, e);
                throw new RuntimeException("groovy script exception", e);
            }
            return new ScriptEngineWrapper(scriptEngine, code);
        });

        return (T) ((Invocable) engineWrapper.getEngine()).invokeFunction(func, args);
    }

    /**
     * 获取编译次数
     */
    public static Long getCompiledCount() {
        return compileCount.get();
    }

    /**
     * 查询已编译的脚本信息
     */
    public static Map<String, String> getCompiledCodes() {
        Map<String, String> codes = new HashMap<>();
        engineMap.forEach((k, v) -> codes.put(k, v.getCode()));
        return codes;
    }
}
